#!/bin/sh

#
# Script which configures balena keys into hardware secure elements and
# prevents updating to a non-signed system if secure boot is enabled
#

set -o errexit

# shellcheck source=/dev/null
. /usr/libexec/os-helpers-logging
# shellcheck disable=SC1091
[ -f /usr/libexec/os-helpers-tpm2 ] && . /usr/libexec/os-helpers-tpm2

DURING_UPDATE="${DURING_UPDATE:-0}"
UMOUNT_EFIVARS=0
SECUREBOOT_VAR="8be4df61-93ca-11d2-aa0d-00e098032b8c-SecureBoot"
SETUPMODE_VAR="8be4df61-93ca-11d2-aa0d-00e098032b8c-SetupMode"
EFIVAR_RE="s,^[^ ]*  *\([^ ]*\) .*$,\1,"
EFIVARS_MOUNTDIR="/sys/firmware/efi/efivars"

exit_with_msg() {
    result=$?
    exit ${result}
}
trap exit_with_msg EXIT

mountEfiVars() {
    # Mount efivarfs if necessary
    if ! mount | grep -q "type efivarfs"; then
        mount -t efivarfs efivarfs "${EFIVARS_MOUNTDIR}"
        UMOUNT_EFIVARS=1
    fi
}

umountEfiVars() {
    # Leave the system in the original state - unmount efivarfs if we mounted it
    if [ "${UMOUNT_EFIVARS}" = 1 ]; then
        umount "${EFIVARS_MOUNTDIR}"
        UMOUNT_EFIVARS=0
    fi
}

updateKeys() {
    SESSION_CTX=$(mktemp -t)
    EFI_MOUNT_DIR="/mnt/efi"
    PASSPHRASE_FILE="$(mktemp -t)"
    CURRENT_POLICY_PATH="$(find /mnt/efi -name "policies.*")"
    DISK_ENC_DIR=""
    for UNLOCK_PCRS in  0,2,3,7 0,1,2,3; do
        tpm2_startauthsession --policy-session -S "${SESSION_CTX}"
        tpm2_policypcr -S "${SESSION_CTX}" -l "sha256:${UNLOCK_PCRS}"
        POLICIES="$(find "${CURRENT_POLICY_PATH}" -type f | sort | xargs)"
        if [ "$(echo "${POLICIES}" | wc -w)" -gt 1 ]; then
            tpm2_policyor -S "${SESSION_CTX}" "sha256:$(echo "${POLICIES}" | sed 's/ /,/g')"
        fi

        if tpm_nvram_retrieve_passphrase "session:${SESSION_CTX}" "${PASSPHRASE_FILE}"; then
            UNLOCK_PCRS_SUCCESS="${UNLOCK_PCRS}"
            info "Retrieved passphrase from TPM NVRAM"
        elif hw_decrypt_passphrase "$EFI_MOUNT_DIR" "session:${SESSION_CTX}" "${PASSPHRASE_FILE}"; then
            UNLOCK_PCRS_SUCCESS="${UNLOCK_PCRS}"
            DISK_ENC_DIR="$(mktemp -d)"
            info "Decrypted passphrase stored on disk"
        fi

        if [ -n "${UNLOCK_PCRS_SUCCESS}" ]; then
            break;
        fi

        tpm2_flushcontext "${SESSION_CTX}" 2>/dev/null
    done

    if [ -n "${UNLOCK_PCRS_SUCCESS}" ]; then
        info "Authorized access to passphrase using PCRs ${UNLOCK_PCRS_SUCCESS}"
    else
        fail "Failed to update policy sealing LUKS passphrase"
    fi

    INACTIVE_SYSROOT=/mnt/sysroot/inactive
    GRUB_BIN="$(find "${INACTIVE_SYSROOT}" -name bootx64.efi.secureboot -print -quit)"
    KERNEL_BIN="$(find "${INACTIVE_SYSROOT}" -name bzImage -print -quit)"
    if [ -z "${GRUB_BIN}" ] || [ -z "${KERNEL_BIN}" ]; then
        fail "Unable to add kernel and bootloader hashes to PCR7 digest"
    fi

    EFI_BINARIES="${GRUB_BIN} ${KERNEL_BIN}"
    PCRS="0,2,3,7"
    PCR_VAL_BIN="$(mktemp -t)"
    PCR_VAL_BIN_EFIBIN="$(mktemp -t)"

    # This OS release may not have the bootloader version required
    # to read the TPM event log, which means we can't assess if the
    # firmware measures EFI binary signatures into PCR7.
    #
    # Create a combined policy that authenticates with PCR7 values
    # calculated with and without the EFI binary hashes.
    POLICY_PATH="$(mktemp -t -d policies.XXXXX)"
    POLICY_PCR="$(mktemp -t)"
    POLICY_EFIBIN="$(mktemp -t)"
    generate_pcr_digests \
        "${PCRS}" \
        "${PCR_VAL_BIN}" \
        "use_pcr2"
    generate_pcr_digests \
        "${PCRS}" \
        "${PCR_VAL_BIN_EFIBIN}" \
        "use_pcr2" \
        "${EFI_BINARIES}"

    tpm2_createpolicy --policy-pcr \
                      -l "sha256:${PCRS}" \
                      -f "${PCR_VAL_BIN}" \
                      -L "${POLICY_PCR}"
    tpm2_createpolicy --policy-pcr \
                      -l "sha256:${PCRS}" \
                      -f "${PCR_VAL_BIN_EFIBIN}" \
                      -L "${POLICY_EFIBIN}"

    case "$(firmware_measures_efibins)" in
        measured)
            info "Using PCR7 digest with EFI binary measurements"
            cp "${POLICY_EFIBIN}" "${POLICY_PATH}/policy.pcr-efibin"
            print_pcr_val_bin "${PCRS}" "${PCR_VAL_BIN_EFIBIN}"
            ;;
        unmeasured)
            info "Using PCR7 digest without EFI binary measurements"
            cp "${POLICY_PCR}" "${POLICY_PATH}/policy.pcr"
            print_pcr_val_bin "${PCRS}" "${PCR_VAL_BIN}"
            ;;
        unknown)
            # we don't have access to the TPM event log, and can't
            # definitively tell whether or not EFI binaries are measured
            # into PCR7, so unlock with both digests
            cp "${POLICY_PCR}" "${POLICY_PATH}/policy.pcr"
            cp "${POLICY_EFIBIN}" "${POLICY_PATH}/policy.pcr-efibin"

            print_pcr_val_bin "${PCRS}" "${PCR_VAL_BIN_EFIBIN}"
            printf "\nOR\n"
            print_pcr_val_bin "${PCRS}" "${PCR_VAL_BIN}"
            ;;
    esac

    tpm_nvram_store_passphrase "${PASSPHRASE_FILE}" "${POLICY_PATH}" "${DISK_ENC_DIR}"
    cp -rf "${POLICY_PATH}" "${EFI_MOUNT_DIR}"
    rm -rf "${CURRENT_POLICY_PATH}"

    # If we are migrating from encrypted passphrase stored in the boot partition
    # to passphrase stored in TPM NVRAM, the above has re-encrypted the passphrase
    # and we now need to store it in the EFI partition, otherwise rollback won't work.
    if [ -n "${DISK_ENC_DIR}" ]; then
        tpm2_evictcontrol -c "${EFI_MOUNT_DIR}/balena-luks.ctx"

        mv "${DISK_ENC_DIR}/persistent.ctx" "${EFI_MOUNT_DIR}/balena-luks.ctx"
        mv "${DISK_ENC_DIR}/passphrase.enc" "${EFI_MOUNT_DIR}/balena-luks.enc"

        rm -rf "${DISK_ENC_DIR}"

        sync "${EFI_MOUNT_DIR}"
    fi

    # update the second stage bootloader kernel binary, as the hash of the
    # original may not be present in the signature database
    rsync --checksum --inplace "${KERNEL_BIN}"{,.sig} "${EFI_MOUNT_DIR}/"

    # update GRUB if necessary
    rsync --checksum --inplace "${GRUB_BIN}" "${EFI_MOUNT_DIR}/EFI/BOOT/bootx64.efi"
    sync

    # This only updates db at this moment
    # PK and KEK need to be implemented
    DB_SYSFS_FILE="db-d719b2cb-3d3a-4596-a3bc-dad00e67656f"
    chattr -i "${EFIVARS_MOUNTDIR}/${DB_SYSFS_FILE}"
    efi-updatevar -a -f "/resin-boot/balena-keys/db.auth" db
    chattr +i "${EFIVARS_MOUNTDIR}/${DB_SYSFS_FILE}"

    # dbx is updated by rollback-health when it confirms that HUP went through
    # here we just copy out the dbx update from the active partition to a pending
    # directory for rollback-health to find it later
    PENDING_DBX_DIR="/mnt/data/balenahup/pending-dbx"
    ACTIVE_ESL_FILE=$(find "/mnt/sysroot/active" -name "db.esl")
    ACTIVE_DBX_FILE=$(find "/mnt/sysroot/active" -name "dbx.auth")
    ACTIVE_DBX_FILE_MD5SUM=$(md5sum "${ACTIVE_DBX_FILE}" | cut -d " " -f 1)
    DEST_FILE="${PENDING_DBX_DIR}/dbx-${ACTIVE_DBX_FILE_MD5SUM}.auth"

    mkdir -p "${PENDING_DBX_DIR}"
    cp -a "${ACTIVE_DBX_FILE}" "${DEST_FILE}"
    sig-list-to-certs "${ACTIVE_ESL_FILE}" "${DEST_FILE}" >/dev/null
}

# Only applicable in the new OS container
if [ "${DURING_UPDATE}" != 1 ]; then
    exit 0
fi

# Only applicable on UEFI systems
if [ ! -d "/sys/firmware/efi" ]; then
    exit 0
fi

mountEfiVars

# Read EFI variables
SECUREBOOT_VAL=$(efivar -p -n "${SECUREBOOT_VAR}" | tail -n 1 | sed -e "${EFIVAR_RE}")
SETUPMODE_VAL=$(efivar -p -n "${SETUPMODE_VAR}" | tail -n 1 | sed -e "${EFIVAR_RE}")

# Only applicable when in secure boot mode
if [ -z "${SECUREBOOT_VAL}" ] || [ "${SECUREBOOT_VAL}" -ne "1" ]; then
    umountEfiVars
    exit 0
fi

# Not applicable when secure boot is in setup mode
if [ -n "${SETUPMODE_VAL}" ] && [ "${SETUPMODE_VAL}" -ne "0" ]; then
    umountEfiVars
    exit 0
fi

# Not applicable if the signatures are in place
if [ -f "/resin-boot/EFI/BOOT/grub-luks.cfg.sig" ]; then
    updateKeys
    umountEfiVars
    exit $?
fi

umountEfiVars

# At this point we are sure the current OS is running in secure boot mode
# and the new image is not signed, abort the update
error "Trying to update to an unsigned OS version while in secure boot mode"
exit 1
